package com.artisan;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.output.Response;

import java.util.concurrent.TimeUnit;

import static dev.langchain4j.data.message.UserMessage.userMessage;
import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;


/**
 * Hello world!
 */
public class App {
    public static void main(String[] args) {
        // 小试牛刀
        firstTalkWitOpenAI();

        // 多轮对话
        multiTalk();
        multiTalk2();
        streamingResponse();
    }

    private static void streamingResponse() {


        StreamingChatLanguageModel model = OpenAiStreamingChatModel.builder()
                .baseUrl("http://langchain4j.dev/demo/openai/v1")
                .modelName(GPT_4_O_MINI)
                .apiKey("demo")
                .build();


        model.generate("tell me a joke", new StreamingResponseHandler<AiMessage>() {
            @Override
            public void onNext(String token) {
                System.out.println(token);
                try {
                    TimeUnit.MILLISECONDS.sleep(200);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }

            }

            @Override
            public void onError(Throwable error) {
                System.out.println("--->:" + error);
            }
        });
    }

    private static void multiTalk() {

        OpenAiChatModel openAiChatModel = OpenAiChatModel.builder()
                .apiKey("demo")
                .modelName(GPT_4_O_MINI)
                .build();

        System.out.println(openAiChatModel.generate("我是小工匠"));
        System.out.println(openAiChatModel.generate("请问，我叫什么名字？"));

        System.out.println("==============");


        UserMessage userMessage1 = userMessage("你好，我是小工匠");
        Response<AiMessage> response1 = openAiChatModel.generate(userMessage1);
        AiMessage aiMessage1 = response1.content(); // 大模型的第一次响应
        System.out.println(aiMessage1.text());
        System.out.println("----");

        // 下面一行代码是重点
        Response<AiMessage> response2 = openAiChatModel.generate(userMessage1, aiMessage1, userMessage("请问，我叫什么？"));
        AiMessage aiMessage2 = response2.content(); // 大模型的第二次响应
        System.out.println(aiMessage2.text());

    }

    /**
     * 演示如何在多轮对话中使用聊天记忆。
     * 该方法展示了如何使用聊天记忆与语言模型进行交互，包括如何向聊天记忆中添加用户消息和AI响应，以及如何控制对话流程。
     */
    private static void multiTalk2() {
        // 初始化聊天记忆，设置最大令牌数为300，并使用OpenAI分词器。
        // 这个设置有助于控制对话长度，避免过多的令牌消耗。
        ChatMemory chatMemory = TokenWindowChatMemory.withMaxTokens(300, new OpenAiTokenizer());

        // 创建并配置语言模型。
        // 使用OpenAI提供的GPT-4-O-MINI模型，并使用演示API密钥。
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey("demo")
                .modelName(GPT_4_O_MINI)
                .build();

        // 你可以完全控制聊天记忆。可以决定是否将特定消息添加到记忆中
        // （例如，你可能不想存储少量样本以节省令牌）。 如果需要，你可以在保存之前处理或修改消息。

        // 向聊天记忆中添加用户消息。
        chatMemory.add(userMessage("你好, 我的名字是小工匠"));
        // 根据当前聊天记忆生成AI响应并打印。
        AiMessage answer = model.generate(chatMemory.messages()).content();
        System.out.println(answer.text());
        // 将AI的响应添加到聊天记忆中。
        chatMemory.add(answer);

        // 向聊天记忆中添加另一个用户消息。
        chatMemory.add(userMessage("我叫什么名字?"));
        // 根据更新后的聊天记忆再次生成AI响应并打印。
        AiMessage answerWithName = model.generate(chatMemory.messages()).content();
        System.out.println(answerWithName.text());
        // 将新的AI响应添加到聊天记忆中。
        chatMemory.add(answerWithName);
    }


    /**
     * 首次与OpenAI进行对话的示例方法
     * 该方法展示了如何使用OpenAI的GPT-4模型进行简单的文本生成
     */
    private static void firstTalkWitOpenAI() {
        // 定义API密钥，此处使用示例密钥"demos"
        String apiKey = "demo";

        // 构建OpenAI聊天模型实例，指定模型名称为GPT_4_O_MINI，并提供API密钥
        ChatLanguageModel model = OpenAiChatModel.builder()
                .modelName(GPT_4_O_MINI)
                .apiKey(apiKey)
                .build();

        // 使用模型生成对问题“你叫什么名字？”的回答，并打印回答内容
        String answer = model.generate("你叫什么名字？");
        System.out.println(answer);

        // 使用模型生成对问题“Say 'Hello World'”的回答，并打印回答内容
        answer = model.generate("Say 'Hello World'");
        System.out.println(answer); // Hello World
    }
}
